import tensorflow as tf
from tensorflow.keras import layers, models
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MultiLabelBinarizer
import pandas as pd

# 加载数据集
movies_df = pd.read_csv('movies.csv')
ratings_df = pd.read_csv('ratings.csv')

# 处理电影类别数据
genres = movies_df['genres'].apply(lambda x: x.split('|'))
mlb = MultiLabelBinarizer()
genres_encoded = mlb.fit_transform(genres)

# 创建电影的特征输入
movie_input = tf.keras.Input(shape=(genres_encoded.shape[1],), name='Movie-Input')

# 创建网络层
movie_embedding = layers.Dense(128, activation='relu')(movie_input)
movie_embedding = layers.Dense(64, activation='relu')(movie_embedding)
movie_embedding = layers.Dense(32, activation='relu')(movie_embedding)

# 输出层
output = layers.Dense(genres_encoded.shape[1], activation='sigmoid')(movie_embedding)

# 创建模型
model = models.Model(inputs=movie_input, outputs=output)
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])

# 训练模型
tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir='./logs', histogram_freq=1)
model.fit(genres_encoded, genres_encoded, epochs=7, callbacks=[tensorboard_callback])

import numpy as np
# 推荐函数
def recommend(movie_id, top_n=5):
    # 我们需要使用np.expand_dims来增加输入数据的维度
    movie_vec = genres_encoded[movie_id]
    movie_vec = np.expand_dims(movie_vec, axis=0)  # 这将数据的形状从(18,)转变为(1, 18)
    prediction = model.predict(movie_vec)
    prediction = np.squeeze(prediction)  # 把结果的维度从(1, 18)降低回(18,)

    # 获取排序后的最高评分电影的索引
    sorted_indexes = prediction.argsort()[::-1][:top_n]

    # 返回对应的电影标题
    recommended_movie_ids = movies_df.iloc[sorted_indexes]['movieId']
    return movies_df[movies_df['movieId'].isin(recommended_movie_ids)]['title']


# import pdb
# pdb.set_trace()
# 推荐与movieId为1的电影相似的5部电影
print(recommend(1))
